#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from retrying import retry
import cStringIO
import pycurl
import os
import sys


def Verify_Meta_XML(filename):
    import commands

    cmd = "grep '</PubmedArticleSet>' " + filename
    result = commands.getoutput(cmd)
    if result.find("</PubmedArticleSet>") != -1:
        return True
    elif result.find("No such file or directory") != -1:
        return False
    else:
        return False


@retry
def Get_PubMed_Meta(filename):
    print("String downloading pubmed meta xml by " + filename + " ... ")
    ids = ""
    infile = open(filename, "r")
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        ids += splitline[0] + ","
    infile.close()
    ids = ids[:-1]
    xml_name = filename[:-3] + "xml"
    if os.path.exists(xml_name):
        os.remove(xml_name)
    url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi'
    buf = cStringIO.StringIO()
    c = pycurl.Curl()
    c.setopt(pycurl.POST, 1)
    c.setopt(c.URL, url)
    c.setopt(pycurl.HTTPPOST, [('db', 'pubmed'), ('id', ids), ('retmode', 'xml')])
    c.setopt(c.WRITEFUNCTION, buf.write)
    c.perform()
    outfile = open(xml_name, "w")
    outfile.write(buf.getvalue())
    outfile.close()
    buf.close()
    c.close()
    print "Verifying " + xml_name + " ... ",
    if not Verify_Meta_XML(xml_name):
        print("invalid. Try to redownload...")
        raise IOError("Invalid xml.")
    else:
        print("OK.")
        print(xml_name + " download complete.")


if __name__ == '__main__':
    Get_PubMed_Meta(sys.argv[1])
